#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_sort_float.h>
#include <gsl/gsl_heapsort.h>
#ifdef USEOPENMP
#include <omp.h>
#endif

#include "halofinder.h"

typedef struct {
  int ind;
  peanokey key;
} SortHaloKey;

static int compSortHaloKey(const void *a, const void *b)
{
  if(((const SortHaloKey*)a)->key > ((const SortHaloKey*)b)->key)
    return 1;
  else if(((const SortHaloKey*)a)->key < ((const SortHaloKey*)b)->key)
    return -1;
  else
    return 0;
}

void halopercolation_serial(void)
{
  int i,j,periodic,NumNbrs,maxNumNbrs,haloInd;
  kdTreeData *td;
  float searchRadius;
  Halo *tmpHalos;
  float minHaloMass;
  float *px,*py,*pz;
  float *nbrsRad2;
  int *nbrsInd;
  float maxMass;
  float dx;
  int N,x,y,z;
  SortHaloKey *sh;
  
  if(ThisTask == 0)
    {
      logProfileTag(PROFILETAG_HALOPERC);
      
#ifdef DEBUG
      fprintf(stderr,"%d: doing serial halo percolation, NumHalos = %ld\n",ThisTask,NumHalos);
#endif
      
      //remove halos with masses that are too small
      minHaloMass = haloFinderData.MinNumPartsHalo;
      haloInd = 0;
      for(i=0;i<NumHalos;++i)
	if(Halos[i].mass >= minHaloMass)
	  {
	    Halos[haloInd] = Halos[i];
	    ++haloInd;
	  }

      if(haloInd < NumHalos)
	{
	  tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*haloInd);
	  assert(tmpHalos != NULL);
	  Halos = tmpHalos;
	  NumHalos = haloInd;
	}
      
      //setup - sort halos by mass, make them into parts for nnbrs finding, use keep vector to mark halos to keep during percolation
#ifdef PERCOLATE_MASS
      sort_halos_mass();
#else
      sort_halos_vcirc();
#endif
      
      N = 1;
      N = N << BITS_PER_DIMENSION;
      dx = (float) (haloFinderData.BoxLengthSnapshot/((double) N));
      sh = ( SortHaloKey*)malloc(sizeof(SortHaloKey)*NumHalos);
      assert(sh != NULL);
      
      for(i=0;i<NumHalos;++i)
	{
	  x = (int) (Halos[i].pos[0]/dx);
	  if(x == N)
	    --x;
	  y = (int) (Halos[i].pos[1]/dx);
	  if(y == N)
	    --y;
	  z = (int) (Halos[i].pos[2]/dx);
	  if(z == N)
	    --z;
	  sh[i].key = peano_hilbert_key(x,y,z,BITS_PER_DIMENSION);
	  sh[i].ind = i;
	}
      
      qsort(sh,(size_t) NumHalos,sizeof(SortHaloKey),compSortHaloKey);
      
      px = (float*)malloc(sizeof(float)*NumHalos);
      assert(px != NULL);
      py = (float*)malloc(sizeof(float)*NumHalos);
      assert(py != NULL);
      pz = (float*)malloc(sizeof(float)*NumHalos);
      assert(pz != NULL);
      for(i=0;i<NumHalos;++i)
	{
	  sh[i].key = 1;
	  px[i] = Halos[sh[i].ind].pos[0];
	  py[i] = Halos[sh[i].ind].pos[1];
	  pz[i] = Halos[sh[i].ind].pos[2];
	}
            
#ifndef LIGHTCONE
      periodic = 1;
#else
      periodic = 0;
#endif
      td = buildkdTree(px,py,pz,(int) NumHalos,NULL);
      
      //set searchRadius by max exclusion distance
      maxMass = Halos[0].mass;
      for(i=0;i<NumHalos;++i)
	if(Halos[i].mass > maxMass)
	  maxMass = Halos[i].mass;
      
      //set max num nbrs
      maxNumNbrs = 100000;
      nbrsRad2 = (float*)malloc(sizeof(float)*maxNumNbrs);
      assert(nbrsRad2 != NULL);
      nbrsInd = (int*)malloc(sizeof(int)*maxNumNbrs);
      assert(nbrsInd != NULL);
      
#ifdef DEBUG
      if(DEBUG_LEVEL > 0)
	fprintf(stderr,"max search radius = %f (%f domain buff size), minHaloMass = %f\n",
		halo_exlcusion_radius(maxMass*haloFinderData.PartMass,maxMass*haloFinderData.PartMass),
		haloFinderData.DomainBuffSize,minHaloMass);
#endif
      
      /*
	keep == 1 means that we keep the halo
	move down list of halos from highest mass to lowest and mark halos to get rid of by setting Parts.index = 0
      */
      for(i=0;i<NumHalos;++i)
	{
	  
#ifdef DEBUG
	  if(i%(NumHalos/20) == 0)
	    fprintf(stderr,"%d of %ld (%0.2f percent) done in halo percolation\n",i,NumHalos,((double) i)/((double) NumHalos)*100.0);
#endif
	  
	  if(sh[i].key)
	    {
	      searchRadius = (float) (halo_exlcusion_radius(Halos[i].mass*haloFinderData.PartMass,maxMass*haloFinderData.PartMass));
	      NumNbrs = get_nnbrs_kdtree(Halos[i].pos,searchRadius,periodic,haloFinderData.BoxLengthSnapshot,
					 &nbrsRad2,&nbrsInd,&maxNumNbrs,px,py,pz,td);
	      
	      if(NumNbrs > 0)
		{
		  for(j=0;j<NumNbrs;++j)
		    {
		      nbrsRad2[j] = (float) (sqrt(nbrsRad2[j]));
		      nbrsInd[j] = (int) (sh[nbrsInd[j]].ind);
		    }
		  
		  //mark less massive nbrs inside exclusion radius as junk
		  for(j=0;j<NumNbrs;++j)
		    if(nbrsRad2[j] < halo_exlcusion_radius(Halos[i].mass*haloFinderData.PartMass,Halos[nbrsInd[j]].mass*haloFinderData.PartMass) 
#ifdef PERCOLATE_MASS
		       && Halos[i].mass >= Halos[nbrsInd[j]].mass
#else
		       && Halos[i].vcirc >= Halos[nbrsInd[j]].vcirc
#endif
		       && nbrsInd[j] != i && sh[nbrsInd[j]].key == 1)
		      {
			sh[nbrsInd[j]].key = 0;
		      }
		}
	    }
	}
      
      //now keep all halos not marked as junk
      haloInd = 0;
      for(i=0;i<NumHalos;++i)
	{
	  if(sh[i].key)
	    {
	      Halos[haloInd] = Halos[i];
	      Halos[haloInd].mass = Halos[haloInd].mass*haloFinderData.PartMass;
	      ++haloInd;
	    }
	}
      
      if(haloInd > 0)
	{
	  tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(haloInd));
	  assert(tmpHalos != NULL);
	  Halos = tmpHalos;
	  NumHalos = haloInd;
	  
	  //set the halo inds
	  for(i=0;i<NumHalos;++i)
	    Halos[i].index = i;
	}
      else
	{
	  free(Halos);
	  NumHalos = 0;
	  Halos = NULL;
	}
      
#ifdef DEBUG
      fprintf(stderr,"%d: done with halo percolation, NumHalos = %ld\n",ThisTask,NumHalos);
#endif
      
      free(sh);
      free(px);
      free(py);
      free(pz);
      free(nbrsRad2);
      free(nbrsInd);
      destroykdTree(td);
      
      logProfileTag(PROFILETAG_HALOPERC);
    }
}

